# yolov8计算器
from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2
import tkinter as tk

# 初始化YOLO模型
model = YOLO("yolov8n.pt")
# 打开指定路径的视频文件
cap = cv2.VideoCapture(0)

# 检查摄像头是否成功打开
if not cap.isOpened():
    raise IOError("Cannot open webcam")

# 获取视频的宽度、高度和帧率
w, h, fps = (
    int(cap.get(x))
    for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)
)

# 定义区域的坐标点
region_points = [(0, 300), (1280, 300), (1280, 360), (0, 360)]

# 初始化物体计数器
counter = object_counter.ObjectCounter()
# 设置物体计数器的参数
counter.set_args(
    view_img=True,  # 是否显示图像
    reg_pts=region_points,  # 物体计数的区域点
    classes_names=model.names,  # 类别名字
    draw_tracks=True,
)  # 是否绘制轨迹

# 当视频文件打开时循环处理每帧视频
while cap.isOpened():
    success, im0 = cap.read()  # 读取一帧视频
    if not success:
        # 如果读取失败或视频处理完成，则退出循环
        print(
            "Video frame is empty or video processing has been successfully completed."
        )
        break

    # 使用YOLO模型进行物体追踪
    tracks = model.track(im0, persist=True, show=False)

    # 在视频帧上执行物体计数并返回帧
    im0 = counter.start_counting(im0, tracks)
    
    # 创建一个窗口逐帧展示
    cv2.imshow('mysqlwindows', im0)

cap.release()
cv2.destroyAllWindows()
